#!/usr/bin/env python3
# Author: Armit
# Create Time: 2024/04/24

import warnings ; warnings.filterwarnings(action='ignore', category=UserWarning)

from typing import Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as M
from torchvision.models.vgg import VGG
from torchvision.models.resnet import ResNet, BasicBlock, Bottleneck


def param_cnt(model:nn.Module) -> int:
  return sum([p.numel() for p in model.parameters() if p.requires_grad])


class PretrainedBackbone(nn.Module):

  def __init__(self, name:str):
    super().__init__()

    self.model = getattr(M, name)(pretrained=True)
    self._out_channels = None

  @property
  def out_channels(self):
    return self._out_channels

class VGGBackbone(PretrainedBackbone):

  def __init__(self, name:str='vgg16'):
    super().__init__(name)

    self.model: VGG
    del self.model.avgpool
    del self.model.classifier

    conv: nn.Conv2d = self.model.features[-3]
    assert isinstance(conv, nn.Conv2d), f'unknown layer type: {type(conv)}'
    self._out_channels = conv.out_channels

  def forward(self, x):
    assert VGG.forward
    self = self.model

    x = self.features(x)
    return x

class ResnetBackbone(PretrainedBackbone):

  def __init__(self, name:str='resnet18'):
    super().__init__(name)

    self.model: ResNet
    del self.model.avgpool
    del self.model.fc

    block = self.model.layer4[-1]
    if isinstance(block, BasicBlock):
      self._out_channels = block.conv2.out_channels
    elif isinstance(block, Bottleneck):
      self._out_channels = block.conv3.out_channels
    else: raise TypeError(f'unknown block type: {type(block)}')


  def forward(self, x):
    assert ResNet.forward
    self = self.model

    x = self.conv1(x)
    x = self.bn1(x)
    x = self.relu(x)
    x = self.maxpool(x)
    x = self.layer1(x)
    x = self.layer2(x)
    x = self.layer3(x)
    x = self.layer4(x)
    return x


class Projector(nn.Module):
  pass

class LinearProjector(Projector):
  
  def __init__(self, d_in:int, d_out:int):
    super().__init__()

    self.fc = nn.Linear(d_in, d_out)

  def forward(self, x):
    return self.fc(x)

class MLPProjector(Projector):

  def __init__(self, d_in:int, d_out:int, dim:int=256):
    super().__init__()

    dim = min(dim, d_in) or max(d_in, d_out)
    self.mlp = nn.Sequential(
      nn.Linear(d_in, dim),
      nn.ReLU(inplace=True),
      nn.Linear(dim, d_out),
    )

  def forward(self, x):
    return self.mlp(x)

class ReLU6MLPProjector(Projector):

  def __init__(self, d_in:int, d_out:int, dim:int=256):
    super().__init__()

    dim = min(dim, d_in) or max(d_in, d_out)
    self.mlp = nn.Sequential(
      nn.Linear(d_in, dim),
      nn.ReLU6(inplace=True),
      nn.Linear(dim, d_out),
    )

  def forward(self, x):
    return self.mlp(x)

class SiLUMLPProjector(Projector):

  def __init__(self, d_in:int, d_out:int, dim:int=256):
    super().__init__()

    dim = min(dim, d_in) or max(d_in, d_out)
    self.mlp = nn.Sequential(
      nn.Linear(d_in, dim),
      nn.SiLU(inplace=True),
      nn.Linear(dim, d_out),
    )

  def forward(self, x):
    return self.mlp(x)


BACKBONE_CLASSES  = {k: v for k, v in globals().items() if type(v) == type(PretrainedBackbone) and issubclass(v, PretrainedBackbone) and v is not PretrainedBackbone}
PROJECTOR_CLASSES = {k[:-len('Projector')]: v for k, v in globals().items() if type(v) == type(Projector)          and issubclass(v, Projector)          and v is not Projector}

def get_backbone(name:str) -> PretrainedBackbone:
  for cls_name, cls in BACKBONE_CLASSES.items():
    prefix = cls_name[:-len('Backbone')].lower()
    if name.startswith(prefix):
      return cls(name)
  raise TypeError(f'no PretrainedBackbone class found for {name}')

def get_projector(name:str, **kwargs) -> Projector:
  for prefix, cls in PROJECTOR_CLASSES.items():
    if name.startswith(prefix):
      return cls(**kwargs)
  raise TypeError(f'no Projector class found for {name}')


class ExtracorModel(nn.Module):

  def __init__(self, backbone:str='vgg16', projector:str='mlp', pool_size:Union[int, Tuple[int, int]]=3, d_out:int=32):
    super().__init__()

    self.backbone = get_backbone(backbone)
    self.avg_pool = nn.AdaptiveAvgPool2d(pool_size)
    d_in = self.backbone.out_channels * (pool_size**2 if isinstance(pool_size, int) else (pool_size[0] * pool_size[1]))
    self.projector = get_projector(projector, d_in=d_in, d_out=d_out)

  def forward(self, x):
    x = self.backbone(x)    # [B, C, H, W]
    x = self.avg_pool(x)    # [B, C, h<H, w<W]
    x = x.flatten(1)        # [B, C*h*w]
    x = self.projector(x)   # [B, D]
    return x


if __name__ == '__main__':
  model = ExtracorModel('vgg16', 'MLP')
  print(model)
  print('param_cnt:', param_cnt(model))

  X = torch.rand([4, 3, 800, 600])
  out = model(X)
  print(out.shape)
